#include <algorithm>
#include <chrono>
#include <cstdlib>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <string>
#include <thread>
#include <vector>
#include <list>
#include <mutex>
#include <atomic>

#include <zstd.h>

namespace fs = std::filesystem;

// use 4 threads to processing file
constexpr int MAX_THREAD = 32;
ZSTD_CCtx *ctxs[MAX_THREAD] {};

struct TaskParameter {
    std::string name;
    fs::path    in;
    fs::path    out;
};

std::list<TaskParameter> task_list;
std::mutex list_mtx;

static void zstd_init()
{
    for (auto &ctx: ctxs) {
        ctx = ZSTD_createCCtx();
    }
}

static void zstd_clean()
{
    for (auto &ctx: ctxs) {
        ZSTD_freeCCtx (ctx);
        ctx = nullptr;
    }
}

std::atomic<bool> running{true};

static std::vector<unsigned char> compress_file(const fs::path &input, size_t &length, ZSTD_CCtx *ctx) {
    std::vector<unsigned char> data;

    if (auto size = fs::file_size (input); size > 0) {
        std::ifstream ifs(input, std::ios::binary | std::ios::in);
        if (ifs.is_open()) {
            std::vector<unsigned char> file_content;
            file_content.resize (size);
            ifs.read ((char*)file_content.data(), size);
            data.resize (ZSTD_compressBound (size));
            length = ZSTD_compressCCtx (ctx, data.data(), data.size(), file_content.data(), file_content.size(), ZSTD_maxCLevel());
        } else {
            length = 0;
        }
    } else {
        length = 0;
    }

    return data;
}

static bool compile_file(const std::string &name, const fs::path &input, const fs::path &output, int ctx_id)
{
    // got var prefix
    std::string var_prefix = name;
    std::ranges::for_each(var_prefix, [](auto &ch){
        if (ch == '/' || ch == '.' || ch == ' ')
            ch = '_';
    });

    // std::cout << name << '{' << input << "to" << output << "}\n";

    // vars
    // const char   {name}_data[]
    // const size_t {name}_length
    // const size_t {name}_origin_length

    // open output file
    std::ofstream ofs(output);
    if (ofs.is_open()) {
        ofs << "// AUTO GENERATED, DO NOT MODIFY!!\n\n";
        ofs << "#include <stddef.h>\n\n";
        ofs << "const char          " << var_prefix << "_name[] = \"/" << name << "\";\n";
        ofs << "const size_t        " << var_prefix << "_origin_length = " << fs::file_size(input) << ";\n";

        size_t length = 0;
        auto const compressed = compress_file (input, length, ctxs[ctx_id]);

        ofs << "const size_t        " << var_prefix << "_length = " << length << ";\n";
        ofs << "const unsigned char " << var_prefix << "_data[] = {\n";
        for (size_t i=0; i<length; ++i)
        {
            if (i%16==0)
                ofs << "\n    ";
            ofs << (int)compressed[i] << ',';
        }
        ofs << "\n};\n";
        return true;
    }
    return false;
}

static void fetch_task_func(int ctx_idx)
{
    while(running) {
        bool feched = false;
        TaskParameter param;
        {
            if (list_mtx.try_lock()) {
                if (!task_list.empty()) {
                    param = task_list.front();
                    task_list.pop_front();
                    feched = true;
                }
                list_mtx.unlock();
            }
        }
        if (feched) {
            if (!compile_file(param.name, param.in, param.out, ctx_idx))
                running = false; // error abort
        }
        std::this_thread::sleep_for(std::chrono::milliseconds(1));
    }
}

static int do_compile(const fs::path &input_list_file) {
    std::ifstream ifs(input_list_file);
    if (ifs.is_open()) {
        std::string root_dir;
        if (std::getline(ifs, root_dir, '\n') && fs::is_directory (root_dir))
        {
            std::string line;
            while(std::getline(ifs, line, '\n'))
            {
                // pattern is rel_path:output_path
                auto const pos = line.find (':');
                if (pos < line.length()) {
                    std::string name = line.substr (0, pos);
                    fs::path in = fs::path{root_dir}/name;
                    fs::path out = fs::path{line.substr (pos+1)};
                    
                    {
                        std::lock_guard lck(list_mtx);
                        task_list.emplace_back(name, in, out);
                    }
                }
            }
            return EXIT_SUCCESS;
        } else {
            std::cerr << "bad root dir: " << root_dir << ", or bad list file\n";
        }
    }
    return EXIT_FAILURE;
}

int main(int argc, char *argv[])
{
    zstd_init();
    // start compile thread
    std::thread *compile_thrs[MAX_THREAD];
    const int use_thread = std::clamp<int>(std::thread::hardware_concurrency()-1, 1, MAX_THREAD);
    for (int i=0; i<use_thread; ++i) {
        compile_thrs[i] = new std::thread(&fetch_task_func, i);
    }
    auto const code = argc == 2 ? do_compile (argv[1]) : 1;

    while (running) {
        // check task
        {
            std::lock_guard lck(list_mtx);
            if (task_list.empty())
                running = false;
        }
        std::this_thread::sleep_for(std::chrono::milliseconds(50));
    }

    // wait for thread done
    for (int i=0; i<use_thread; ++i) {
        auto &t = compile_thrs[i];
        if (t->joinable())
            t->join();
        delete t;
    }
    zstd_clean();
    return task_list.empty() ? code : EXIT_FAILURE;
}
